import datetime
from typing import List, Tuple, Dict, Any, Union

import adata
import baostock as bs
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib import gridspec
from tqdm import tqdm

from backTest import BacktestEngine
from specificStockAnalysis import SupportResistanceAnalyzer  # 导入你的分析器类
from concurrent.futures import ProcessPoolExecutor, as_completed

# 设置字体支持中文
plt.rcParams['font.sans-serif'] = ['SimHei']  # 'SimHei' 是黑体，你可以根据系统安装的字体更改
plt.rcParams['axes.unicode_minus'] = False  # 正确显示负号

# 配置参数
INITIAL_CAPITAL = 100000.0  # 初始资金
RISK_FREE_RATE = 0.03       # 无风险利率（年化）
MIN_DATA_LENGTH = 20        # 信号计算所需最小数据长度


def calculate_volatility(df, annualize=True):
    """
    计算股票的平均波动率（年化或日波动率）

    参数:
    df (pd.DataFrame): 包含收盘价的 DataFrame，列名为 'Close'
    annualize (bool): 是否年化波动率，默认为 True

    返回:
    float: 波动率（年化或日波动率）
    """
    # 计算每日收益率
    df['Daily Return'] = df['close'].pct_change()

    # 删除缺失值（第一天的收益率无法计算）
    df.dropna(inplace=True)

    # 计算日波动率（收益率的标准差）
    daily_volatility = df['Daily Return'].std()

    # 年化波动率（假设一年有 252 个交易日）
    if annualize:
        annualized_volatility = daily_volatility * np.sqrt(252)
        return annualized_volatility
    else:
        return daily_volatility


def get_past_date(n_days):
    current_time = datetime.datetime.now()
    past_date = current_time - datetime.timedelta(days=n_days)
    return past_date.strftime('%Y-%m-%d')

def load_data(data) -> pd.DataFrame:
    """加载股票数据，确保格式正确"""
    # trade_date 列转成datetime格式并设置为索引
    data['trade_date'] = pd.to_datetime(data['trade_date'])
    data.set_index('trade_date', inplace=True)
    return data[['open', 'high', 'low', 'close', 'volume']]
def optimize_parameters(data, rr_range: List[float], stop_loss_pct: float = 0.02) -> pd.DataFrame:
    """多参数优化"""
    data = load_data(data)
    engine = BacktestEngine(data)
    results = []
    trades = []
    portforios = []

    for rr in rr_range:
        result, trade, portfolio = engine.run_backtest(risk_reward_ratio=rr,stop_loss_pct = stop_loss_pct)
        results.append(result)
        trades.append(trade)
        portforios.append(portfolio)

    df = pd.DataFrame(results)
    return df

def run_backtest_with_params(stock_code, start_date, rr_range):
    res_df = adata.stock.market.get_market(stock_code=stock_code, k_type=1, start_date=start_date)
    yearly_return = calculate_volatility(res_df, False)


    """辅助函数，用于多线程调用"""

    result = optimize_parameters(res_df, rr_range, stop_loss_pct=yearly_return)
    return result


if __name__ == "__main__":
    # 示例参数范围
    rr_range = np.arange(0.8, 2.0, 0.4).tolist()

    adata.proxy(is_proxy=True, ip='127.0.0.1:10810')
    # 登陆系统
    # lg = bs.login()
    # # 显示登陆返回信息
    # print('login respond error_code:' + lg.error_code)
    # print('login respond  error_msg:' + lg.error_msg)

    # 获取股票代码和名称
    # rs = bs.query_all_stock()  # 选择一个日期，这里选择2023-10-01
    # data_list = []
    # while (rs.error_code == '0') & rs.next():
    #     data_list.append(rs.get_row_data())

    res_df = adata.stock.info.all_code()
    # 组合 exchange 和 stock_code 成 "SZ.00001" 形式的列表
    # 创建字典，将exchange和stock_code组合成key，short_name作为value
    stock_code_name_map = {f"{row['exchange']}.{row['stock_code']}": row['short_name'] for index, row in res_df.iterrows()}

    # # 将数据转换为DataFrame
    # stock_df = pd.DataFrame(data_list, columns=rs.fields)
    #
    # # 构建股票代码到股票名称的字典
    # stock_code_name_map = {row['code']: row['code_name'] for index, row in stock_df.iterrows()}

    # 示例：获取3天前的日期
    n = 400
    start_date = get_past_date(n)
    count = 0


    colume_list = ['stock_code','stock_name','risk_reward_ratio','sharpe_ratio','total_return','max_drawdown','num_trades','win_rate','volatility']
    total_df = pd.DataFrame(columns=colume_list)


    for c_stock_code, name in stock_code_name_map.items():
        try:
            stock_code = c_stock_code.split(".")[1]
            if "ST" in name:
                continue
            res_df = adata.stock.market.get_market(stock_code=stock_code, k_type=1, start_date=start_date)
            try:
                info_df = adata.stock.info.get_industry_sw(stock_code=stock_code)
            except:
                pass
            if res_df.empty:
                continue

            res_df = adata.stock.market.get_market(stock_code=stock_code, k_type=1, start_date=start_date)

            yearly_return = calculate_volatility(res_df, False)
            # 运行优化
            results_df = optimize_parameters(res_df, rr_range, yearly_return)

            # 输出最佳参数
            best_row = results_df.loc[results_df['sharpe_ratio'].idxmax()]
            best_row['stock_code'] = stock_code
            best_row['stock_name'] = name
            best_row['volatility'] = yearly_return
            total_df = pd.concat([total_df, best_row.to_frame().T], ignore_index=True)
            count += 1
            if count % 10 == 0:
                # 输出该结果到csv
                total_df.to_csv("optimization_results_{}.csv".format(count), index=False)

            # 更新后的输出格式
            print(f"""
            ================ {name}优化结果 ================
            统计次数：       {count}
            最佳盈亏比:      {best_row['risk_reward_ratio']:.2f}
            夏普比率:       {best_row['sharpe_ratio']:.2f}
            总收益率:       {best_row['total_return']:.2f}%
            最大回撤:       {best_row['max_drawdown']:.2f}%
            胜率:          {best_row['win_rate']:.1f}%
            交易次数:       {best_row['num_trades']}
            日波动率：       {yearly_return:.2f}%
            ========================================
            """)
        except Exception as e:
            print(f"Error occurred for stock {c_stock_code}: {e}")

    total_df.to_csv('optimization_results_total.csv', index=False)
